#ifndef JIT_COMPILER_H
#define JIT_COMPILER_H

#include "../../include/gudb.h"
#include "../memory/memory_manager.h"
#include <stdint.h>
#include <stdbool.h>
#include <ctype.h>
#include <pthread.h>

// JIT Configuration
#define JIT_MAX_EXPRESSIONS         1024    // Maximum cached expressions
#define JIT_COMPILATION_THRESHOLD   100     // Compile after N interpretations
#define JIT_CODE_CACHE_SIZE         (4 * 1024 * 1024)  // 4MB code cache
#define JIT_MAX_PARAMETERS          32      // Maximum expression parameters
#define JIT_MAX_INSTRUCTION_SIZE    16      // Maximum instruction size in bytes
#define JIT_OPTIMIZATION_LEVELS     4       // Number of optimization levels

// Forward declarations
typedef struct jit_compiler jit_compiler_t;
typedef struct jit_expression jit_expression_t;
typedef struct jit_context jit_context_t;
typedef struct jit_value jit_value_t;
typedef struct jit_instruction jit_instruction_t;
typedef struct jit_basic_block jit_basic_block_t;
typedef struct jit_function jit_function_t;

// JIT Data Types
typedef enum {
    JIT_TYPE_INT8,
    JIT_TYPE_INT16,
    JIT_TYPE_INT32,
    JIT_TYPE_INT64,
    JIT_TYPE_UINT8,
    JIT_TYPE_UINT16,
    JIT_TYPE_UINT32,
    JIT_TYPE_UINT64,
    JIT_TYPE_FLOAT,
    JIT_TYPE_DOUBLE,
    JIT_TYPE_BOOL,
    JIT_TYPE_STRING,
    JIT_TYPE_TIMESTAMP,
    JIT_TYPE_DECIMAL,
    JIT_TYPE_NULL,
    JIT_TYPE_UNKNOWN
} jit_type_t;

// JIT Operations
typedef enum {
    // Arithmetic operations
    JIT_OP_ADD,
    JIT_OP_SUB,
    JIT_OP_MUL,
    JIT_OP_DIV,
    JIT_OP_MOD,
    JIT_OP_NEG,
    
    // Comparison operations
    JIT_OP_EQ,
    JIT_OP_NE,
    JIT_OP_LT,
    JIT_OP_LE,
    JIT_OP_GT,
    JIT_OP_GE,
    
    // Logical operations
    JIT_OP_AND,
    JIT_OP_OR,
    JIT_OP_NOT,
    
    // Bitwise operations
    JIT_OP_BIT_AND,
    JIT_OP_BIT_OR,
    JIT_OP_BIT_XOR,
    JIT_OP_BIT_NOT,
    JIT_OP_LSHIFT,
    JIT_OP_RSHIFT,
    
    // String operations
    JIT_OP_CONCAT,
    JIT_OP_SUBSTR,
    JIT_OP_LENGTH,
    JIT_OP_LIKE,
    JIT_OP_ILIKE,
    
    // Math functions
    JIT_OP_ABS,
    JIT_OP_SQRT,
    JIT_OP_POWER,
    JIT_OP_LOG,
    JIT_OP_EXP,
    JIT_OP_SIN,
    JIT_OP_COS,
    JIT_OP_TAN,
    
    // Aggregate functions
    JIT_OP_SUM,
    JIT_OP_AVG,
    JIT_OP_COUNT,
    JIT_OP_MIN,
    JIT_OP_MAX,
    
    // Control flow
    JIT_OP_CASE,
    JIT_OP_WHEN,
    JIT_OP_IF,
    JIT_OP_COALESCE,
    
    // Memory operations
    JIT_OP_LOAD,
    JIT_OP_STORE,
    JIT_OP_MOVE,
    JIT_OP_CAST,
    
    // Function calls
    JIT_OP_CALL,
    JIT_OP_RETURN,
    
    // Constants and variables
    JIT_OP_CONST,
    JIT_OP_PARAM,
    JIT_OP_VAR
} jit_op_t;

// JIT Value representation
struct jit_value {
    jit_type_t type;
    uint32_t id;                    // Unique value ID
    bool is_constant;               // Whether value is compile-time constant
    bool is_parameter;              // Whether value is runtime parameter
    
    union {
        int64_t int_val;
        uint64_t uint_val;
        double float_val;
        bool bool_val;
        char* string_val;
        void* ptr_val;
    } data;
    
    // Register assignment (for compiled code)
    int8_t reg_id;                  // -1 if not in register
    bool spilled;                   // Whether value is spilled to memory
    uint32_t spill_offset;          // Stack offset if spilled
};

// JIT Instruction
struct jit_instruction {
    jit_op_t opcode;
    jit_value_t* dest;              // Destination value
    jit_value_t* src1;              // First source operand
    jit_value_t* src2;              // Second source operand
    jit_value_t* src3;              // Third source operand (for ternary ops)
    
    uint32_t id;                    // Instruction ID
    uint8_t* machine_code;          // Generated machine code
    size_t machine_code_size;       // Size of machine code
    
    struct jit_instruction* next;   // Next instruction in block
    struct jit_instruction* prev;   // Previous instruction in block
};

// JIT Basic Block
struct jit_basic_block {
    uint32_t id;                    // Block ID
    jit_instruction_t* first_inst;  // First instruction
    jit_instruction_t* last_inst;   // Last instruction
    uint32_t inst_count;            // Instruction count
    
    // Control flow
    jit_basic_block_t* successor1;   // First successor block
    jit_basic_block_t* successor2;   // Second successor block (for branches)
    jit_basic_block_t** predecessors; // Array of predecessor blocks
    uint32_t predecessor_count;     // Number of predecessors
    
    // Code generation
    uint8_t* machine_code;          // Generated machine code
    size_t machine_code_size;       // Size of machine code
    size_t machine_code_offset;     // Offset in function code
    
    // Analysis data
    bool visited;                   // For traversal algorithms
    uint32_t dom_depth;             // Dominator tree depth
    jit_basic_block_t* dom_parent;  // Immediate dominator
    
    struct jit_basic_block* next;   // Next block in function
};

// JIT Function
struct jit_function {
    char name[64];                  // Function name
    uint32_t id;                    // Function ID
    
    // Parameters and return
    jit_type_t return_type;         // Return value type
    jit_type_t* param_types;        // Parameter types
    uint32_t param_count;           // Number of parameters
    
    // Basic blocks
    jit_basic_block_t* entry_block; // Entry basic block
    jit_basic_block_t* blocks;      // Linked list of all blocks
    uint32_t block_count;           // Number of blocks
    
    // Values
    jit_value_t** values;           // Array of all values
    uint32_t value_count;           // Number of values
    uint32_t next_value_id;         // Next available value ID
    
    // Compiled code
    uint8_t* machine_code;          // Generated machine code
    size_t machine_code_size;       // Size of machine code
    void* compiled_function;        // Function pointer
    
    // Optimization data
    uint32_t execution_count;       // Times executed
    uint64_t total_execution_time;  // Total execution time in nanoseconds
    uint32_t optimization_level;    // Current optimization level
    bool is_compiled;               // Whether function is compiled
    
    // Memory management
    size_t stack_size;              // Required stack size
    uint32_t register_usage;        // Number of registers used
    
    pthread_mutex_t mutex;          // Function-level synchronization
};

// JIT Expression (high-level representation)
struct jit_expression {
    char source[256];               // Original expression string
    uint64_t hash;                  // Expression hash for lookup
    
    jit_function_t* function;       // Compiled function
    
    // Execution statistics
    uint32_t execution_count;       // Number of times executed
    uint64_t total_time;            // Total execution time
    uint64_t compilation_time;      // Time spent compiling
    uint32_t cache_hits;            // Result cache hits
    uint32_t cache_misses;          // Result cache misses
    
    // Performance thresholds
    bool should_compile;            // Whether to compile this expression
    bool is_hot;                    // Whether expression is frequently used
    uint32_t complexity_score;      // Expression complexity estimate
    
    // Result caching
    struct {
        void* cached_result;        // Cached result value
        uint64_t cache_key;         // Cache key
        bool valid;                 // Whether cache is valid
        uint64_t timestamp;         // Cache timestamp
    } result_cache;
    
    // Linked list for hash table
    struct jit_expression* next;
    struct jit_expression* prev;
};

// JIT Compilation Context
struct jit_context {
    jit_function_t* current_function; // Function being built
    jit_basic_block_t* current_block;  // Block being built
    
    // Value management
    jit_value_t** value_stack;      // Value stack for expression building
    uint32_t stack_top;             // Stack top pointer
    uint32_t stack_capacity;        // Stack capacity
    
    // Error handling
    char error_message[256];        // Last error message
    int error_code;                 // Last error code
    
    // Temporary storage
    uint8_t* temp_buffer;           // Temporary code buffer
    size_t temp_buffer_size;        // Size of temporary buffer
    size_t temp_buffer_pos;         // Current position in buffer
};

// JIT Compiler main structure
struct jit_compiler {
    // Expression cache
    jit_expression_t** expression_cache; // Hash table of expressions
    uint32_t cache_size;            // Size of hash table
    uint32_t cached_expressions;    // Number of cached expressions
    
    // Code memory management
    uint8_t* code_memory;           // Executable memory pool
    size_t code_memory_size;        // Size of code memory
    size_t code_memory_used;        // Used code memory
    pthread_mutex_t code_mutex;     // Code memory synchronization
    
    // Compilation settings
    uint32_t optimization_level;    // Global optimization level
    bool enable_profiling;          // Whether to enable profiling
    bool enable_caching;            // Whether to enable result caching
    uint32_t compilation_threshold; // Interpretation count before compilation
    
    // Performance monitoring
    uint64_t total_compilations;    // Total number of compilations
    uint64_t total_compilation_time; // Total time spent compiling
    uint64_t total_executions;      // Total number of executions
    uint64_t cache_hits;            // Total cache hits
    uint64_t cache_misses;          // Total cache misses
    
    // Statistics per operation type
    struct {
        uint64_t count;             // Number of operations
        uint64_t total_time;        // Total execution time
        uint32_t avg_time;          // Average execution time
    } op_stats[64];                 // Indexed by jit_op_t
    
    // Thread safety
    pthread_rwlock_t cache_lock;    // Cache read-write lock
    pthread_mutex_t stats_mutex;    // Statistics mutex
    
    // Background optimization
    pthread_t optimizer_thread;     // Background optimizer thread
    bool optimizer_running;         // Whether optimizer is running
    pthread_cond_t optimizer_cond;  // Optimizer condition variable
    pthread_mutex_t optimizer_mutex; // Optimizer mutex
    
    // Memory manager
    memory_manager_t* memory_mgr;   // Memory manager for allocations
};

// JIT Compiler API

// Compiler management
jit_compiler_t* jit_compiler_create(memory_manager_t* memory_mgr);
void jit_compiler_destroy(jit_compiler_t* compiler);
int jit_compiler_set_optimization_level(jit_compiler_t* compiler, uint32_t level);

// Expression compilation and execution
jit_expression_t* jit_compile_expression(jit_compiler_t* compiler, const char* expression);
jit_value_t* jit_execute_expression(jit_expression_t* expr, jit_value_t** parameters);
void jit_expression_destroy(jit_expression_t* expr);

// Function building API
jit_function_t* jit_function_create(jit_compiler_t* compiler, const char* name, 
                                   jit_type_t return_type, jit_type_t* param_types, 
                                   uint32_t param_count);
void jit_function_destroy(jit_function_t* function);
int jit_function_compile(jit_function_t* function);

// Basic block management
jit_basic_block_t* jit_block_create(jit_function_t* function);
void jit_block_set_current(jit_context_t* ctx, jit_basic_block_t* block);
int jit_block_add_successor(jit_basic_block_t* block, jit_basic_block_t* successor);

// Value management
jit_value_t* jit_value_create_const_int(jit_context_t* ctx, int64_t value);
jit_value_t* jit_value_create_const_float(jit_context_t* ctx, double value);
jit_value_t* jit_value_create_const_string(jit_context_t* ctx, const char* value);
jit_value_t* jit_value_create_parameter(jit_context_t* ctx, uint32_t param_index, jit_type_t type);
jit_value_t* jit_value_create_temp(jit_context_t* ctx, jit_type_t type);

// Instruction generation
jit_instruction_t* jit_emit_binary_op(jit_context_t* ctx, jit_op_t op, 
                                     jit_value_t* left, jit_value_t* right);
jit_instruction_t* jit_emit_unary_op(jit_context_t* ctx, jit_op_t op, jit_value_t* operand);
jit_instruction_t* jit_emit_call(jit_context_t* ctx, void* function_ptr, 
                                jit_value_t** args, uint32_t arg_count);
jit_instruction_t* jit_emit_return(jit_context_t* ctx, jit_value_t* value);
jit_instruction_t* jit_emit_branch(jit_context_t* ctx, jit_value_t* condition, 
                                  jit_basic_block_t* true_block, jit_basic_block_t* false_block);

// Type system
jit_type_t jit_type_from_string(const char* type_name);
const char* jit_type_to_string(jit_type_t type);
size_t jit_type_size(jit_type_t type);
bool jit_type_is_numeric(jit_type_t type);
bool jit_type_is_integer(jit_type_t type);
jit_type_t jit_type_promote(jit_type_t type1, jit_type_t type2);

// Optimization passes
int jit_optimize_function(jit_function_t* function, uint32_t optimization_level);
int jit_optimize_constant_folding(jit_function_t* function);
int jit_optimize_dead_code_elimination(jit_function_t* function);
int jit_optimize_common_subexpression_elimination(jit_function_t* function);
int jit_optimize_loop_invariant_code_motion(jit_function_t* function);
int jit_optimize_register_allocation(jit_function_t* function);

// Code generation
int jit_generate_machine_code(jit_function_t* function);
int jit_generate_x86_64(jit_function_t* function);
int jit_generate_arm64(jit_function_t* function);

// Performance monitoring
void jit_compiler_print_stats(const jit_compiler_t* compiler);
void jit_expression_print_stats(const jit_expression_t* expr);
double jit_compiler_get_cache_hit_ratio(const jit_compiler_t* compiler);
uint64_t jit_compiler_get_average_compilation_time(const jit_compiler_t* compiler);

// Adaptive optimization
int jit_adaptive_optimize(jit_compiler_t* compiler);
bool jit_should_compile_expression(const jit_expression_t* expr);
uint32_t jit_estimate_expression_complexity(const char* expression);
void* jit_optimizer_thread_func(void* arg);

// Cache management
uint64_t jit_hash_expression(const char* expression);
jit_expression_t* jit_cache_lookup(jit_compiler_t* compiler, const char* expression);
int jit_cache_insert(jit_compiler_t* compiler, jit_expression_t* expr);
int jit_cache_evict_lru(jit_compiler_t* compiler);

// Memory management
int jit_allocate_executable_memory(jit_compiler_t* compiler, size_t size, void** ptr);
int jit_free_executable_memory(jit_compiler_t* compiler, void* ptr, size_t size);
int jit_make_memory_executable(void* ptr, size_t size);

// Context management
jit_context_t* jit_context_create(jit_compiler_t* compiler);
void jit_context_destroy(jit_context_t* ctx);
int jit_context_push_value(jit_context_t* ctx, jit_value_t* value);
jit_value_t* jit_context_pop_value(jit_context_t* ctx);

// Debugging and introspection
void jit_function_print_ir(const jit_function_t* function);
void jit_function_print_assembly(const jit_function_t* function);
int jit_function_validate(const jit_function_t* function);
int jit_dump_machine_code(const jit_function_t* function, const char* filename);

#endif // JIT_COMPILER_H